import AsyncHTTPClient
import NIO
import NIOSSL
import NIOHTTP1

class MyResponseAccumulator: HTTPClientResponseDelegate {
    public struct Response {
        /// Remote host of the request.
        public var host: String
        /// Response HTTP status.
        public var status: HTTPResponseStatus
        /// Reponse HTTP headers.
        public var headers: HTTPHeaders
        /// Response body.
        public var body: ByteBuffer?
    }
    
    enum State {
        case idle
        case head(HTTPResponseHead)
        case body(HTTPResponseHead, ByteBuffer)
        case end
        case error(Error)
    }
    
    var state = State.idle
    let request: HTTPClient.Request
    
    var head: HTTPResponseHead? = nil
    var body: ByteBuffer? = nil
    
    init(request: HTTPClient.Request) {
        self.request = request
    }
    
    func didReceiveHead(task: HTTPClient.Task<Response>, _ head: HTTPResponseHead) -> EventLoopFuture<Void> {
        switch self.state {
        case .idle:
            self.state = .head(head)
        case .head:
            preconditionFailure("head already set")
        case .body:
            preconditionFailure("no head received before body")
        case .end:
            preconditionFailure("request already processed")
        case .error:
            break
        }
        return task.eventLoop.makeSucceededFuture(())
    }
    
    func didReceivePart(task: HTTPClient.Task<Response>, _ part: ByteBuffer) -> EventLoopFuture<Void> {
        switch self.state {
        case .idle:
            preconditionFailure("no head received before body")
        case .head(let head):
            self.state = .body(head, part)
            self.head = head
            self.body = part
        case .body(let head, var body):
            var part = part
            body.writeBuffer(&part)
            self.state = .body(head, body)
            self.body = body
        case .end:
            preconditionFailure("request already processed")
        case .error:
            break
        }
        return task.eventLoop.makeSucceededFuture(())
    }
    
    func didReceiveError(task: HTTPClient.Task<Response>, _ error: Error) {
        self.state = .error(error)
    }
    
    func didFinishRequest(task: HTTPClient.Task<Response>) throws -> Response {
        switch self.state {
        case .idle:
            preconditionFailure("no head received before end")
        case .head(let head):
            return Response(host: self.request.host, status: head.status, headers: head.headers, body: nil)
        case .body(let head, let body):
            return Response(host: self.request.host, status: head.status, headers: head.headers, body: body)
        case .end:
            preconditionFailure("request already processed")
        case .error(let error):
            switch error {
            case NIOSSLError.uncleanShutdown:
                if let head = self.head, (!head.headers.contains(name: "Content-Length") && head.headers["Transfer-Encoding"].first != "chunked") {
                    if let body = self.body {
                        return Response(host: self.request.host, status: head.status, headers: head.headers, body: body)
                    } else {
                        return Response(host: self.request.host, status: head.status, headers: head.headers, body: nil)
                    }
                } else {
                    throw error
                }
            default:
                throw error
            }
        }
    }
}

let eventLoop = MultiThreadedEventLoopGroup(numberOfThreads: 1).next()
let httpClient = HTTPClient(eventLoopGroupProvider: .shared(eventLoop))

let request = try HTTPClient.Request(url: "https://accounts.google.com/o/oauth2/token", method: .POST, body: .string(""))

let delegate = MyResponseAccumulator(request: request)
let task = httpClient.execute(request: request, delegate: delegate)
try task.futureResult.flatMap { response in
    return eventLoop.makeSucceededFuture(response)
}.flatMapError { error -> EventLoopFuture<MyResponseAccumulator.Response> in
    do {
        let response = try delegate.didFinishRequest(task: task)
        return eventLoop.makeSucceededFuture(response)
    } catch {
        return eventLoop.makeFailedFuture(error)
    }
}.flatMap { response -> EventLoopFuture<()> in
    print(response.status.code)
    if var responseBody = response.body {
        print(responseBody.readString(length: responseBody.readableBytes) ?? "")
    }
    return eventLoop.makeSucceededFuture(())
}.wait()